#根据框切图

import glob
import os
import shutil
import xml.etree.cElementTree as et
import cv2
import fire
import numpy as np
import sys



def get_dots_by_xml(xml_file, sw_mode):  
    '''
    从XML文件中获取所有角点坐标，返回以左上角角点和右下角角点组成的点集的集合
    xml_path:单个xml文件
    '''
    tree = et.parse(xml_file)  # 从xml文件读取数据
    root = tree.getroot()  # 加载xml得到根节点
    dots = {}

    name_sw = {
        sw_mode[0]:'pbp',
        sw_mode[1]:'pbbp',
        sw_mode[2]:'cbp',
        sw_mode[3]:'cbbp'
    }
    
    for cls in ['pbp', 'pbbp', 'cbp', 'cbbp']:
        dots[cls] = []
    for obj in root.findall('object'):  # 获得所有对象框
        name = obj.find('name').text  # 获取节点\对象框名
        if sw_mode[-1]:
            if name in name_sw:
                name = name_sw[name]  # 支持wqsa快捷键
        bbox = obj.find('bndbox')  # 查找bnbox，找到四个角点
        bbox = [int(bbox.find(tag).text) for tag in ['xmin', 'ymin', 'xmax', 'ymax']]  # 获取四个角点坐标
        start_dot = np.array(bbox[:2])  # 找出左上角角点坐标
        end_dot = np.array(bbox[2:])  # 找出右下角角点坐标
        dots[name].append((start_dot, end_dot))  # 把点加入dots
    return dots



def generate_points_dict(xml_dir, sw_mode):
    '''
    以一张图片的所有标签生成的点集生成字典
    xml_dir:保存所有xml的路径
    '''
    xml_list = glob.glob(os.path.join(xml_dir, '*.xml'))  # 获得xml_dir路径下的xml文件
    dots_dict = {}
    for xml_file in xml_list:
        dots = get_dots_by_xml(xml_file, sw_mode)  # 遍历每个xml文件的所有标签
        dots_dict[xml_file] = dots
    return dots_dict  # 返回字典



def img_cut(img_file, dots, save_path, cls, slice):
    '''
    切图保存
    img_file:单张图片文件
    dots:这张图片对应的点集
    save_path:保存图片的根路径
    cls:类别名称，包括pbp、pbbp、cbp、cbbp（片并片、片不并片、串并片、串不并片）
    slice:控制切分段图与整图的方法，False为切整图
    '''
    print('img {} is processing'.format(os.path.basename(img_file).split('.')[0]))
    tmp_img = cv2.imdecode(np.fromfile(img_file, dtype=np.uint8), 1)
    i = 0
    for one_cls in cls:  # 分别遍历不同类型，把不同类型的小图保存到不同路径
        # print(one_cls, dots[one_cls])
        rot, size = check_cls(one_cls, slice)
        # img_save_path = os.path.join(save_path, one_cls, 'img')
        img_save_path = os.path.join(save_path, one_cls)
        for points in dots[one_cls]:
            start_point = points[0]
            end_point = points[1]
            # print('left_up corner: {}; right_down corner: {}'.format(start_point, end_point))
            roi = tmp_img[start_point[1]:end_point[1], start_point[0]:end_point[0]]
            if rot:
                fin_img = np.rot90(roi)
            else:
                fin_img = roi
            if slice:
                BLACK = [0, 0, 0]
                fin_img = cv2.resize(fin_img, size)
                if rot:
                    fin_img = cv2.copyMakeBorder(fin_img, 295, 295, 0, 0, cv2.BORDER_CONSTANT, value = BLACK)
                else:
                    fin_img = cv2.copyMakeBorder(fin_img, 562, 563, 0, 0, cv2.BORDER_CONSTANT, value = BLACK)
            # else:
            #     fin_img = cv2.resize(fin_img, size)
            img_name = ''.join([os.path.basename(img_file).split('.')[0], '-{}.jpg'.format(i)])
            save_img = os.path.join(img_save_path, img_name)
            cv2.imwrite(save_img, fin_img)
            i += 1



# def img_resize(img_file, size, tmp_img_path, dot_dict):
#     '''
#     调整大小
#     '''

#     img1 = cv2.imdecode(np.fromfile(img_file, dtype=np.uint8), 1)
#     img2 = cv2.resize(img1, size)
#     new_img_file = os.path.join(tmp_img_path, os.path.basename(img_file))
#     cv2.imwrite(new_img_file, img2)
#     y1, x1 = img1.shape[:2]
#     y2, x2 = img2.shape[:2]
#     xrate = x2 / x1
#     yrate = y2 / y1
#     point_li = []
#     for point in dot_dict:
#         new_point = (int(point[0] * xrate), int(point[1] * yrate))
#         point_li.append(new_point)
#     return new_img_file, point_li



def check_dir(dir_path):
    '''
    检测路径，如果有路径提示是否删除，无路径新建路径
    '''
    if os.path.exists(dir_path):
        print('!!!dir {} exists'.format(dir_path))
        comfirm = input("!!!remove dir, yes or mix? [y or m]: ")
        if comfirm == 'y':
            print('!!!{} have been removed!!!'.format(dir_path))
            shutil.rmtree(dir_path)
            os.makedirs(dir_path)
        elif comfirm != 'm':
            print('did not remove!')
            sys.exit(0)
    else:
        os.makedirs(dir_path)



def check_cls(one_cls, slice):
    '''
    传入类别，返回大小、是否旋转
    one_cls:类别名称，包括pbp、pbbp、cbp、cbbp（片并片、片不并片、串并片、串不并片）
    slice:控制切分段图与整图的方法，False为切整图
    '''
    if slice:
        if one_cls == 'pbp':
            size = (750, 160)
            rot = True
        elif one_cls == 'pbbp':
            size = (750, 160)
            rot = True
        elif one_cls == 'cbp':
            size = (375, 375)
            rot = False
        elif one_cls == 'cbbp':
            size = (375, 375)
            rot = False
    else:
        if one_cls == 'pbp':
            size = (750, 750)
            rot = True
        elif one_cls == 'pbbp':
            size = (750, 750)
            rot = True
        elif one_cls == 'cbp':
            size = (375, 1500)
            rot = False
        elif one_cls == 'cbbp':
            size = (375, 1500)
            rot = False
    return rot, size



def each_count(save_path):
    for img_save_path in os.listdir(save_path):
        img_each_path = os.path.join(save_path, img_save_path)
        img_num = len(glob.glob(os.path.join(img_each_path, '*.jpg')))
        print('class {} has {} files!'.format(img_save_path, img_num))



def main(path, save_path, sw_mode = ['w', 'q', 's', 'a', True], slice = False, dir = True):
    '''
    根据传入的path切图，cbp/cbbp/pbp/pbbp对应串并片/串不并片/片并片/片不并片
    path:包含原图img和xml的路径
    save_path:想要保存切图小图的根路径
    sw_mode:四种标签（pbp、pbbp、cbp、cbbp）的快捷键，最后用True/False控制开关，True为开
    slice:控制切分段图与整图的方法，False为切整图
    dir:是否新建路径，True新建路径
    '''
    img_path = os.path.join(path, 'img')
    xml_path = os.path.join(path, 'xml')
    cls = 'pbp,pbbp,cbp,cbbp'
    cls_li = cls.split(',')
    print(cls_li)

    if dir:
        check_path = []
        for one_cls in cls_li:
            # img_save_path = os.path.join(save_path, one_cls, 'img')
            # check_path.append(img_save_path)
            # xml_save_path = os.path.join(save_path, one_cls, 'xml')
            # check_path.append(xml_save_path)
            
            img_save_path = os.path.join(save_path, one_cls)
            check_path.append(img_save_path)
        for dir in check_path:
            check_dir(dir)
       
    dots_dict = generate_points_dict(xml_path, sw_mode)
    xml_list = glob.glob(os.path.join(xml_path, '*.xml'))
    for xml_file in xml_list:
        img_file = os.path.join(img_path, '.'.join([os.path.basename(xml_file).split('.')[0], 'jpg']))
        img_cut(img_file, dots_dict[xml_file], save_path, cls_li, slice)
    
    each_count(save_path)

if __name__ == "__main__":
    fire.Fire(main)
    